import torch
from torch import nn
from torch import optim
from PIL import Image
from torchvision import transforms
import cv2
import os
import numpy as np
os.environ["CUDA_VISIBLE_DEVICES"] = '0'

def interface(img_path, dst_path):
    model = torch.load(
        './CHenTao/front_model-regular-80-0.6821497296026636.pth', map_location='cpu')
    img = Image.open(img_path)
    img = img.resize((512, 512))
    filename = img_path.split('/')[-1]
    filename = filename.split('.')[0]
    print(filename)
    # simple_transform = transforms.ToTensor()
    resize_transform = transforms.Compose([transforms.Resize([512, 512]), transforms.ToTensor()])
    img = resize_transform(img).view([1, 3, 512, 512])

    output1 = model.module(img)
    output1 = output1.squeeze().detach().cpu().numpy()
    output2 = np.zeros((512, 512), dtype=np.uint8)
    temp = np.zeros((512, 512), dtype=np.uint8)
    for j in range(512):
        for k in range(512):

            temp[j, k] = (output1[0, j, k] * 255 + output1[1, j, k] * 255) / 2
            if output1[0, j, k] * 255 > 50:
                output2[j, k] = 128

            if output1[1, j, k] * 255 > 200:
                output2[j, k] = 255

    print(output2.shape)
    # outname = 'result/04ct/color/' + filename + '.tif'
    # cv2.imwrite(outname, temp)
    # outname = 'result/04ct/mask/' + filename + '.tif'
    cv2.imwrite(dst_path, output2)

    return output2